import torch
import torch.multiprocessing as mp
def worker_fn(rank, world_size,a):
    # 在每个进程中执行的函数
    # rank 表示进程的编号，world_size 表示总的进程数量

    # 在这里可以执行训练或推断任务
    # 通常，每个进程会负责不同的数据子集
    # print(f"Now is epoch {world_size} on GPU {rank}")
    a.append(world_size+rank)

if __name__ == "__main__":
    # 设置进程数量
    num_processes = 5
    for epoch in range(10):
        a = mp.Manager().list()
        # 启动多进程
        mp.spawn(
            worker_fn,  # 要在每个进程中运行的函数
            args=(epoch,a),  # 传递给 worker_fn 的参数
            nprocs=num_processes,  # 进程数量
            join=True  # 等待所有进程完成
        )
        print(len(a))
        assert len(a) == 5
        print(sum(a))
